import numpy as np


class LightDataset:

    def divide_list(self, lst: list, n):
        datas = []
        # 使用列表推导式来创建分组
        for i in range(0, len(lst), n):
            if i+n <= len(lst):
                datas.append(lst[i:i + n])
        return datas

    def __init__(self, data_type: str = 'train'):
        # 判断用户选择的数据类型是否是训练集、测试集、验证集三者之一。
        if data_type not in ["train", "test", "valid"]:
            raise ValueError("类型只允许train test valid三种")
        X = []
        y = []
        # 类别绑定
        category_name_dict = {
            "dark": 0,  # 黑暗
            "indirect": 1,  # 侧光
            "room": 2,  # 室内朝上
            "sun": 3  # 阳光直射
        }
        for category in category_name_dict.keys():
            # 同时读取温湿度数据集
            with open("./datasets/{}/light_{}.txt".format(data_type, category), "r") as f:
                data = f.read().strip().split("\n")
                # 按照每十个一组生成返回列表
                data = self.divide_list(data, 10)
                # print(data)
                X = X + data

                y = y + np.full((len(data), 1), category_name_dict[category]).tolist()

        self.X = np.array(X, dtype=np.float32)
        self.y = np.array(y, dtype=np.float32)
        # print(self.X,self.y)
        # print(self.X.shape,self.y.shape)

    def __len__(self):
        return len(self.X)

    def get_data(self):
        return self.X,self.y

    def get_category_name(self):
        return ["黑暗", "室内侧光", "室内照明", "光线直射"]

if __name__ == '__main__':
    dataset = LightDataset()
